#!/bin/bash

# 安装MuJoco200和mujoco-py 2.0.x的自动化脚本

# 检查是否已安装必要工具
check_dependencies() {
    if ! command -v wget &> /dev/null; then
        echo "请先安装wget"
        exit 1
    fi
    
    if ! command -v unzip &> /dev/null; then
        echo "请先安装unzip"
        exit 1
    fi
    
    if ! command -v pip &> /dev/null; then
        echo "请先安装pip"
        exit 1
    fi
}

# 创建.mujoco目录并下载文件
setup_mujoco() {
    # 获取用户主目录
    HOME_DIR=$HOME
    MUJOCO_DIR="$HOME_DIR/.mujoco"
    
    echo "创建.mujoco目录..."
    mkdir -p "$MUJOCO_DIR"
    
    echo "下载mujoco200_linux.zip..."
    wget -q https://www.roboti.us/download/mujoco200_linux.zip -O "$MUJOCO_DIR/mujoco200_linux.zip"
    
    echo "下载mjkey.txt..."
    wget -q https://www.roboti.us/file/mjkey.txt -O "$MUJOCO_DIR/mjkey.txt"
    
    echo "解压mujoco200_linux.zip..."
    unzip -q -o "$MUJOCO_DIR/mujoco200_linux.zip" -d "$MUJOCO_DIR"
    
    # 确保解压后的文件夹名为mujoco200
    if [ -d "$MUJOCO_DIR/mujoco200_linux" ]; then
        mv "$MUJOCO_DIR/mujoco200_linux" "$MUJOCO_DIR/mujoco200"
    fi
    
    # 复制key文件到指定目录
    echo "复制key文件..."
    cp "$MUJOCO_DIR/mjkey.txt" "$MUJOCO_DIR/mujoco200/bin/"
}

# 设置环境变量
setup_environment() {
    HOME_DIR=$HOME
    MUJOCO_PATH="$HOME_DIR/.mujoco/mujoco200/bin"
    
    echo "设置环境变量..."
    
    # 检查环境变量是否已设置，避免重复添加
    if ! grep -q "LD_LIBRARY_PATH.*$MUJOCO_PATH" "$HOME_DIR/.bashrc"; then
        echo "export LD_LIBRARY_PATH=\$LD_LIBRARY_PATH:$MUJOCO_PATH" >> "$HOME_DIR/.bashrc"
    fi
    
    if ! grep -q "MUJOCO_KEY_PATH" "$HOME_DIR/.bashrc"; then
        echo "export MUJOCO_KEY_PATH=~/.mujoco\${MUJOCO_KEY_PATH}" >> "$HOME_DIR/.bashrc"
    fi
    
    # 立即生效环境变量
    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$MUJOCO_PATH
    export MUJOCO_KEY_PATH=~/.mujoco${MUJOCO_KEY_PATH}
    export MUJOCO_GL=egl
    source ~/.bashrc
}

# 安装mujoco-py
install_mujoco_py() {
    echo "安装mujoco-py==2.0.x..."

    # 更新包列表并安装所需依赖
    echo "正在更新包列表..."
    apt-get update

    echo "正在安装libosmesa6-dev..."
    apt-get install -y libosmesa6-dev

    # 安装额外可能需要的依赖（针对无头服务器环境）
    echo "安装额外的OpenGL依赖..."
    apt-get install -y libgl1-mesa-dev libgl1-mesa-glx libglew-dev

    sudo apt-get install patchelf

    cd ~/miniconda3/lib/  # 根据实际路径修改  anaconda3/envs/drl/bin/../lib/
    mv libstdc++.so.6 libstdc++.so.6.old
    ln -s /usr/lib/x86_64-linux-gnu/libstdc++.so.6 libstdc++.so.6

    # pip install Cython>=0.27.2
    pip install 'cython<3'
    pip install glfw
    pip install mujoco_py==2.0.2.5
}

# 主程序
main() {
    check_dependencies
    setup_mujoco
    setup_environment
    install_mujoco_py
    
    echo "安装完成！请重启终端或运行 'source ~/.bashrc' 使环境变量生效"
}

# 执行主程序
main
